import torch
import numpy as np
from sklearn.metrics import classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns

def evaluate_model(model, test_loader, class_names, device=None):
    """
    评估模型性能
    
    参数:
        model (nn.Module): 要评估的模型
        test_loader (DataLoader): 测试数据加载器
        class_names (list): 类别名称列表
        device (torch.device): 设备(CPU/GPU)
        
    返回:
        accuracy (float): 模型在测试集上的准确率
        report (str): 分类报告
        cm (numpy.ndarray): 混淆矩阵
    """
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    model.to(device)
    model.eval()
    
    correct = 0
    total = 0
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
            
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    # 计算准确率
    accuracy = 100. * correct / total
    
    # 生成分类报告
    report = classification_report(
        all_labels, all_preds, 
        target_names=class_names, 
        digits=4
    )
    
    # 生成混淆矩阵
    cm = confusion_matrix(all_labels, all_preds)
    
    # 绘制混淆矩阵
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=class_names, yticklabels=class_names)
    plt.xlabel('Predicted')
    plt.ylabel('Actual')
    plt.title('Confusion Matrix')
    plt.tight_layout()
    plt.savefig('confusion_matrix.png')
    plt.close()
    
    print(f'Test Accuracy: {accuracy:.2f}%')
    print('\nClassification Report:')
    print(report)
    
    return accuracy, report, cm